package cn.leaf.elasticsearch.util;

import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.support.master.AcknowledgedResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.client.indices.CreateIndexResponse;
import org.elasticsearch.client.indices.GetIndexRequest;
import org.elasticsearch.common.text.Text;
import org.elasticsearch.common.unit.Fuzziness;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.FuzzyQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * description: es操作工具类
 *
 * @author Cyril
 * @version v1.0.0
 * @since 2021-05-22 12:39:43
 */
@Component
public class EsOpsUtil<T> {

    private final RestHighLevelClient client;

    public EsOpsUtil(@Qualifier("restHighLevelClient") RestHighLevelClient client) {
        this.client = client;
    }

    /**
     * 判断索引是否存在
     *
     * @param index 索引名称
     * @return boolean
     */
    public boolean existsIndex(String index) throws IOException {
        GetIndexRequest request = new GetIndexRequest(index);
        return client.indices().exists(request, RequestOptions.DEFAULT);
    }

    /**
     * 创建索引
     *
     * @param index 索引名称
     * @return boolean
     */
    public boolean createIndex(String index) throws IOException {
        CreateIndexRequest request = new CreateIndexRequest(index);
        CreateIndexResponse response = client.indices().create(request, RequestOptions.DEFAULT);
        return response.isAcknowledged();
    }

    /**
     * 删除索引
     *
     * @param index 索引名称
     * @return boolean
     */
    public boolean deleteIndex(String index) throws IOException {
        DeleteIndexRequest deleteIndexRequest = new DeleteIndexRequest(index);
        AcknowledgedResponse response = client.indices().delete(deleteIndexRequest, RequestOptions.DEFAULT);
        return response.isAcknowledged();

    }

    /**
     * 判断某索引下文档id是否存在
     *
     * @param id    文档id
     * @param index 索引名称
     * @return boolean
     */
    public boolean existsDoc(String index, String id) throws IOException {
        GetRequest request = new GetRequest(index, id);
        request.fetchSourceContext(new FetchSourceContext(false));
        request.storedFields("_none_");
        return client.exists(request, RequestOptions.DEFAULT);
    }

    /**
     * 添加文档记录
     *
     * @param index 索引名称
     * @param id    文档id
     * @param t     范型
     * @return boolean
     */
    public boolean addDoc(String index, String id, T t) throws IOException {
        IndexRequest request = new IndexRequest(index);
        request.id(id);
        request.timeout(TimeValue.timeValueSeconds(1));
        request.source(JacksonUtil.bean2Json(t), XContentType.JSON);
        IndexResponse response = client.index(request, RequestOptions.DEFAULT);
        RestStatus status = response.status();
        return status == RestStatus.OK || status == RestStatus.CREATED;
    }


    /**
     * 批量添加文档记录
     *
     * @param index 索引名称
     * @param list  文档集合
     * @return boolean
     */
    public boolean bulkAdd(String index, List<T> list) throws IOException {
        BulkRequest bulkRequest = new BulkRequest();
        list.parallelStream().forEach(e ->
                bulkRequest.add(new IndexRequest(index).source(JacksonUtil.bean2Json(e)))
        );
        BulkResponse bulkResponse = client.bulk(bulkRequest, RequestOptions.DEFAULT);
        return !bulkResponse.hasFailures();

    }

    /**
     * 更新文档记录
     *
     * @param index 索引名称
     * @param id    文档id
     * @param t     需要更新的对象
     * @return boolean
     */
    public boolean updateDoc(String index, String id, T t) throws IOException {
        UpdateRequest request = new UpdateRequest(index, id);
        request.doc(JacksonUtil.bean2Json(t));
        request.timeout(TimeValue.timeValueSeconds(1));
        UpdateResponse updateResponse = client.update(
                request, RequestOptions.DEFAULT);
        return updateResponse.status() == RestStatus.OK;
    }

    /**
     * 根据id来获取文档记录
     *
     * @param index 索引名称
     * @param id    文档id
     * @return {@link GetResponse}
     */
    public GetResponse getDoc(String index, String id) throws IOException {
        GetRequest request = new GetRequest(index, id);
        return client.get(request, RequestOptions.DEFAULT);
    }

    /**
     * 删除文档记录
     *
     * @param index 索引名称
     * @param id    文档id
     * @return boolean
     */
    public boolean deleteDoc(String index, String id) throws IOException {
        DeleteRequest request = new DeleteRequest(index, id);
        //timeout
        request.timeout(TimeValue.timeValueSeconds(1));
        DeleteResponse deleteResponse = client.delete(request, RequestOptions.DEFAULT);
        return deleteResponse.status() == RestStatus.OK;
    }

    /**
     * 批量删除文档记录
     *
     * @param index 索引名称
     * @param list  文档id集合
     * @return boolean
     */
    public boolean bulkDelete(String index, List<String> list) throws IOException {
        BulkRequest bulkRequest = new BulkRequest();
        bulkRequest.timeout(TimeValue.timeValueMinutes(2));
        list.parallelStream().forEach(e ->
                bulkRequest.add(new DeleteRequest(index).id(e))
        );
        BulkResponse bulkResponse = client.bulk(bulkRequest, RequestOptions.DEFAULT);
        return !bulkResponse.hasFailures();

    }

    /**
     * 根据某个字段查询
     *
     * @param index    索引名称
     * @param field    字段名称
     * @param key      字段内容
     * @param pageNo   当前页码
     * @param pageSize 每页条数
     * @return {@link List<SearchHit>}
     */
    public List<SearchHit> search(String index, String field, String key, Integer pageNo, Integer pageSize) throws IOException {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        sourceBuilder.query(QueryBuilders.termQuery(field, key));
        //分页公式：(currentPageNo - 1) * pageSize
        int from = (pageNo - 1) * pageSize;
        sourceBuilder.from(from);
        sourceBuilder.size(pageSize);
        searchRequest.source(sourceBuilder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        SearchHits hits = searchResponse.getHits();
        return Arrays.asList(hits.getHits());
    }

    /**
     * 全部搜索
     *
     * @param index    索引名称
     * @param pageNo   当前页码
     * @param pageSize 每页条数
     * @param field1   字段1
     * @param field2   字段2
     * @return {@link List<SearchHit>}
     */
    public List<SearchHit> searchAll(String index, String field1, String field2, Integer pageNo, Integer pageSize) throws IOException {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder builder = new SearchSourceBuilder();
        builder.query(QueryBuilders.matchAllQuery());
        // 根据某个字段排序
        builder.sort(field1, SortOrder.ASC);
        // 过滤结果字段
        String[] include = {field1};
        String[] exclude = {field2};
        builder.fetchSource(include, exclude);
        //分页公式：(currentPageNo - 1) * pageSize
        int from = (pageNo - 1) * pageSize;
        builder.from(from);
        builder.size(pageSize);
        searchRequest.source(builder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        SearchHits hits = searchResponse.getHits();
        return Arrays.asList(hits.getHits());
    }

    /**
     * 组合查询
     *
     * @param index    索引名称
     * @param field    查询字段
     * @param key      查询关键字
     * @param pageNo   当前页码
     * @param pageSize 每页条数
     * @return {@link List<SearchHit>}
     */
    public List<SearchHit> searchBool(String index, String field, String key, Integer pageNo, Integer pageSize) throws IOException {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder builder = new SearchSourceBuilder();
        BoolQueryBuilder boolQuery = QueryBuilders.boolQuery();
        boolQuery.must(QueryBuilders.matchQuery(field, key));
        boolQuery.mustNot(QueryBuilders.matchQuery(field, key));
        boolQuery.should(QueryBuilders.matchQuery(field, key));
        builder.query(boolQuery);
        // 分页公式：(currentPageNo - 1) * pageSize
        int from = (pageNo - 1) * pageSize;
        builder.from(from);
        builder.size(pageSize);
        searchRequest.source(builder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        SearchHits hits = searchResponse.getHits();
        return Arrays.asList(hits.getHits());
    }

    /**
     * 范围查询
     *
     * @param index    索引名称
     * @param field    字段名称
     * @param range1   查询范围
     * @param range2   查询范围
     * @param pageNo   当前页码
     * @param pageSize 每页条数
     * @return {@link List<SearchHit>}
     */
    public List<SearchHit> searchRange(String index, String field, String range1, String range2, Integer pageNo, Integer pageSize) throws IOException {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder builder = new SearchSourceBuilder();
        RangeQueryBuilder rangeQuery = QueryBuilders.rangeQuery(field);
        // gt lt 不包含值 gte lte 包含值
        rangeQuery.gt(range1);
        rangeQuery.lt(range2);
        builder.query(rangeQuery);
        // 分页公式：(currentPageNo - 1) * pageSize
        int from = (pageNo - 1) * pageSize;
        builder.from(from);
        builder.size(pageSize);
        searchRequest.source(builder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        SearchHits hits = searchResponse.getHits();
        return Arrays.asList(hits.getHits());
    }

    /**
     * 模糊查询
     *
     * @param index    索引名称
     * @param field    字段名称
     * @param key      查询关键字
     * @param pageNo   当前页码
     * @param pageSize 每页条数
     * @return {@link List<SearchHit>}
     */
    public List<SearchHit> searchFuzziness(String index, String field, String key, Integer pageNo, Integer pageSize) throws IOException {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder builder = new SearchSourceBuilder();
        // 模糊查询 Fuzziness：查询偏移量 相差几个字符查询结果
        FuzzyQueryBuilder fuzzyQueryBuilder = QueryBuilders.fuzzyQuery(field, key).fuzziness(Fuzziness.AUTO);
        builder.query(fuzzyQueryBuilder);
        // 分页公式：(currentPageNo - 1) * pageSize
        int from = (pageNo - 1) * pageSize;
        builder.from(from);
        builder.size(pageSize);
        searchRequest.source(builder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        SearchHits hits = searchResponse.getHits();
        return Arrays.asList(hits.getHits());
    }

    /**
     * 高亮查询
     *
     * @param index    索引名称
     * @param field    字段名称
     * @param key      查询关键字
     * @param pageNo   当前页码
     * @param pageSize 每页条数
     * @return {@link List<SearchHit>}
     */
    public List<Map<String, Object>> searchHighLight(String index, String field, String key, Integer pageNo, Integer pageSize) throws IOException {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder builder = new SearchSourceBuilder();
        // 模糊查询 Fuzziness：查询偏移量 相差几个字符查询结果
        FuzzyQueryBuilder fuzzyQueryBuilder = QueryBuilders.fuzzyQuery(field, key).fuzziness(Fuzziness.ONE);
        builder.query(fuzzyQueryBuilder);

        // 构建高亮字段和属性
        HighlightBuilder highlightBuilder = new HighlightBuilder();
        highlightBuilder.field(field);
        highlightBuilder.preTags("<font color='blue'>");
        highlightBuilder.postTags("</font>");
        builder.highlighter(highlightBuilder);

        // 分页公式：(currentPageNo - 1) * pageSize
        int from = (pageNo - 1) * pageSize;
        builder.from(from);
        builder.size(pageSize);
        searchRequest.source(builder);

        // 查询数据
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        SearchHits hits = searchResponse.getHits();
        List<SearchHit> hitList = Arrays.asList(hits.getHits());

        // 替换高亮内容
        return hitList.parallelStream().map(new Function<SearchHit, Map<String, Object>>() {
            @Override
            public Map<String, Object> apply(SearchHit hit) {
                Map<String, HighlightField> highlightFields = hit.getHighlightFields();
                HighlightField highlightField = highlightFields.get(field);
                Map<String, Object> sourceMap = hit.getSourceAsMap();
                if (Optional.ofNullable(highlightField).isPresent()) {
                    Text[] fragments = highlightField.fragments();
                    StringBuilder fieldValue = new StringBuilder();
                    for (Text text : fragments) {
                        fieldValue.append(text);
                    }
                    sourceMap.put(field, fieldValue.toString());
                }
                return sourceMap;
            }
        }).collect(Collectors.toList());
    }

    /**
     * 获取最大值
     *
     * @param index 索引名称
     * @param field 字段名称
     * @return {@link Aggregations}
     */
    public Aggregations searchMax(String index, String field) throws IOException {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder builder = new SearchSourceBuilder();
        AggregationBuilder aggregationBuilder = AggregationBuilders.max("max-" + field).field(field);
        builder.aggregation(aggregationBuilder);
        searchRequest.source(builder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        return searchResponse.getAggregations();
    }

    /**
     * 获取最小值
     *
     * @param index 索引名称
     * @param field 字段名称
     * @return {@link Aggregations}
     */
    public Aggregations searchMin(String index, String field) throws IOException {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder builder = new SearchSourceBuilder();
        AggregationBuilder aggregationBuilder = AggregationBuilders.min("max-" + field).field(field);
        builder.aggregation(aggregationBuilder);
        searchRequest.source(builder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        return searchResponse.getAggregations();
    }

    /**
     * 根据字段分组
     *
     * @param index 索引名称
     * @param field 字段名称
     * @return {@link Aggregations}
     */
    public Aggregations searchGroup(String index, String field) throws IOException {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder builder = new SearchSourceBuilder();
        AggregationBuilder aggregationBuilder = AggregationBuilders.terms(field + "-group").field(field);
        builder.aggregation(aggregationBuilder);
        searchRequest.source(builder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        return searchResponse.getAggregations();
    }
}